#-*- coding:utf-8 -*-
# author:贤宁
# datetime:2021/11/5 9:43
# software: PyCharm

import torch
from torch import nn
import torch.nn.functional as F
from torch.nn.utils import spectral_norm
from utils import get_conv_layer
from torch.distributions import uniform
"""三种残差块"""
#G Block
class G_Block(nn.Module):
    def __init__(self,in_channel,out_channel):
        super(G_Block,self).__init__()
        self.block1 = nn.Sequential(
            nn.UpsamplingBilinear2d(scale_factor=2),  #线性上采样
            nn.Conv2d(in_channel, out_channel, 1),   #1x1 卷积
        )

        self.block2 = nn.Sequential(
            nn.BatchNorm2d(in_channel),
            nn.LeakyReLU(0.2,inplace=True),
            nn.UpsamplingBilinear2d(scale_factor=2),   #线性上采样
            nn.Conv2d(in_channel, in_channel, 3, stride=1, padding=1),   #这里特征图与通道数都不变
            nn.BatchNorm2d(in_channel),
            nn.LeakyReLU(0.2,inplace=True),
            nn.Conv2d(in_channel, out_channel, 3, stride=1, padding=1)   #这里特征图大小不变，通道数增加
        )
    def forward(self,x):
        return self.block1(x) + self.block2(x)   #残差求和
#D&3D Block
class D_Block(nn.Module):
    def __init__(self,in_channel,out_channel,conv_type,first_relu=True,keep_same_output=False):
        """keep_same_output：输出是否与输入具有相同的空间维度，若为FALSE，则缩小为2
           conv_type: 决定卷积的类型,"standard"--->Conv2d,"3d"---->Conv3d
           first_relu:也为一个bool型
        """
        super(D_Block, self).__init__()
        self.keep_same_output = keep_same_output
        self.conv_type = conv_type
        self.first_relu = first_relu
        conv2d = get_conv_layer(conv_type)
        self.conv_1x1 = conv2d(
            in_channels=in_channel,
            out_channels=out_channel,
            kernel_size=1,
        )
        self.first_conv_3x3 = conv2d(
            in_channels=in_channel,
            out_channels=in_channel,
            kernel_size=3,
            padding=1,
        )
        self.last_conv_3x3 = conv2d(
            in_channels=in_channel,
            out_channels=out_channel,
            kernel_size=3,
            padding=1,
            stride=1,
        )
        if conv_type == "3d":
            # 需要谱的归一化
            self.conv_1x1 = spectral_norm(self.conv_1x1)
            self.first_conv_3x3 = spectral_norm(self.first_conv_3x3)
            self.last_conv_3x3 = spectral_norm(self.last_conv_3x3)
        self.relu = torch.nn.LeakyReLU(0.2,inplace=True)
    def forward(self,x):
        x1 = self.conv_1x1(x)
        if not self.keep_same_output:
            x1 = F.interpolate(
                x1, mode="trilinear" if self.conv_type == "3d" else "bilinear", scale_factor=0.5
            )
        if self.first_relu:
            x = self.relu(x)
        x = self.first_conv_3x3(x)
        x = self.relu(x)
        x = self.last_conv_3x3(x)
        if not self.keep_same_output:
            x = F.interpolate(
                x, mode="trilinear" if self.conv_type == "3d" else "bilinear", scale_factor=0.5
            )
        x = x1 + x  # 计算残差
        return x
#L Block(整体均不会改变特征图的大小)
class L_Block(nn.Module):
    def __init__(self,in_channel,out_channel, conv_type="standard"):
        super(L_Block, self).__init__()
        self.conv_type = conv_type
        conv2d = get_conv_layer(conv_type)
        self.conv_1x1 = conv2d(
            in_channels=in_channel,
            out_channels=out_channel - in_channel,
            kernel_size=1,
        )

        self.first_conv_3x3 = conv2d(
            in_channels= in_channel, out_channels=out_channel, kernel_size=3, padding=1, stride=1
        )
        self.relu = torch.nn.LeakyReLU(0.2,inplace=True)
        self.last_conv_3x3 = conv2d(
            in_channels=out_channel,
            out_channels=out_channel,
            kernel_size=3,
            padding=1,
            stride=1,
        )
    def forward(self,x):
        x1 = self.conv_1x1(x)
        #print(x1.size())
        x = self.relu(x)
        x2 = self.first_conv_3x3(x)
        x2 = self.relu(x2)
        x2 = self.last_conv_3x3(x2)
        #print(x2.size())
        x = x2 + (torch.cat((x,x1),dim=1))
        return x
"""Block---End"""

"""潜在条件堆栈:通过将正态分布N(0,1)--->zsp(作为条件输入到G中）"""
class SelfAttention2d(nn.Module):

    def __init__(self, input_dims, output_dims=None, return_attn=False):
        output_dims = input_dims // 8 if output_dims is None else output_dims
        if output_dims == 0:
            raise Exception(
                "The output dims corresponding to the input dims is 0. Increase the input\
                            dims to 8 or more. Else specify output_dims"
            )
        super(SelfAttention2d, self).__init__()
        self.query = nn.Conv2d(input_dims, output_dims, 1)
        self.key = nn.Conv2d(input_dims, output_dims, 1)
        self.value = nn.Conv2d(input_dims, input_dims, 1)
        self.gamma = nn.Parameter(torch.zeros(1))
        self.return_attn = return_attn

    def forward(self, x):
        r"""Computes the output of the Self Attention Layer
        Args:
            x (torch.Tensor): A 4D Tensor with the channel dimension same as ``input_dims``.
        Returns:
            A tuple of the ``output`` and the ``attention`` if ``return_attn`` is set to ``True``
            else just the ``output`` tensor.
        """
        dims = (x.size(0), -1, x.size(2) * x.size(3))
        out_query = self.query(x).view(dims)
        out_key = self.key(x).view(dims).permute(0, 2, 1)
        attn = F.softmax(torch.bmm(out_key, out_query), dim=-1)
        out_value = self.value(x).view(dims)
        out_value = torch.bmm(out_value, attn).view(x.size())
        out = self.gamma * out_value + x
        if self.return_attn:
            return out, attn
        return out

class Latent_Condition_Stack(nn.Module):
    def __init__(self,
        shape,
        output_channels,
        use_attention):
        """
        :param shape: 输入的规格(8,8,8)
        :param out_channel:768
        :param use_attention:为原文中的空间注意模块设定，为bool，默认为True
        """
        super(Latent_Condition_Stack, self).__init__()
        self.use_attention = use_attention
        self.shape = shape
        self.distribution = uniform.Uniform(low=torch.Tensor([0.0]), high=torch.Tensor([1.0]))
        """
        uniform.Uniform官方文档中是随机生成[low,high]之间张量样本，后面的sample(shape)是取出一个shape规格
        的随机张量样本。
        """
        self.conv_3x3 = nn.Conv2d(
            in_channels=shape[0], out_channels=shape[0], kernel_size=3, padding=1
        )
        self.l_block1 = L_Block(shape[0], output_channels // 32)
        self.l_block2 = L_Block(
            output_channels // 32, output_channels // 16
        )
        self.l_block3 = L_Block(
            output_channels // 16, output_channels // 4
        )
        if self.use_attention:
            self.att_block = SelfAttention2d(
                output_channels // 4, output_channels // 4
            )
        self.l_block4 = L_Block(
            output_channels // 4, output_channels
        )
    def forward(self,x):
        z = self.distribution.sample(self.shape)
        z = torch.Tensor.permute(z, (3, 0, 1, 2)).type_as(x)
        z = self.conv_3x3(z)
        z = self.l_block1(z)
        #print(z.size())
        z = self.l_block2(z)
        z = self.l_block3(z)
        if self.use_attention:
            z = self.att_block(z)
        z = self.l_block4(z)
        return z
"""Latent_Condition_Stack---End"""

"""生成器的调节堆栈"""
class Context_Condition_Stack(nn.Module):
    def __init__(self,
        input_channels,
        output_channels,
        num_context_steps,
        conv_type= "standard"):
        """
        :param input_channels:每一步所输入的通道数(1
        :param output_channels:最低块的输出通道数(768
        :param num_context_steps:(4
        :param conv_type:
        """
        super().__init__()
        self.conv_type = conv_type
        self.input_channels = input_channels
        self.output_channels = output_channels
        self.num_context_steps = num_context_steps
        conv2d = get_conv_layer(conv_type)
        #self.S2D = nn.PixelShuffle(upscale_factor=int(0.5))
        #四个下采样的剩余块
        self.d1 = D_Block(
            4 * input_channels,
            ((output_channels // 4) * input_channels) // num_context_steps,
            conv_type,
        )
        self.d2 = D_Block(
            ((output_channels // 4) * input_channels) // num_context_steps,
            ((output_channels // 2) * input_channels) // num_context_steps,
            conv_type,
        )
        self.d3 = D_Block(
            ((output_channels // 2) * input_channels) // num_context_steps,
            (output_channels * input_channels) // num_context_steps,
            conv_type,
        )
        self.d4 = D_Block(
            (output_channels * input_channels) // num_context_steps,
            (output_channels * 2 * input_channels) // num_context_steps,
            conv_type,
        )
        #再将每个d的输出用3x3频谱归一化卷积
        self.conv1 = spectral_norm(
            conv2d(
                in_channels=(output_channels // 4) * input_channels,
                out_channels=(output_channels // 8) * input_channels,
                kernel_size=3,
                padding=1,
            )
        )

        self.conv2 = spectral_norm(
            conv2d(
                in_channels=(output_channels // 2) * input_channels,
                out_channels=(output_channels // 4) * input_channels,
                kernel_size=3,
                padding=1,
            )
        )

        self.conv3 = spectral_norm(
            conv2d(
                in_channels=output_channels * input_channels,
                out_channels=(output_channels // 2) * input_channels,
                kernel_size=3,
                padding=1,
            )
        )

        self.conv4 = spectral_norm(
            conv2d(
                in_channels=output_channels * 2 * input_channels,
                out_channels=output_channels * input_channels,
                kernel_size=3,
                padding=1,
            )
        )

        self.relu = torch.nn.LeakyReLU(0.2,inplace=True)
    def forward(self,x):
        #x = self.S2D(x)
        #print(x.size())
        x = F.interpolate(
                x,  size=[4,128,128]
            )
        steps = x.size(1)  # 步数
        scale_1 = []
        scale_2 = []
        scale_3 = []
        scale_4 = []
        """
        scale来存储每次剩余块的输出，以及将其作为采样堆栈的一部分输入
        """
        for i in range(steps):
            #print(x[:, i, :, :,:].size())
            s1 = self.d1(x[:, i, :, :,:])
            #print(s1.size())
            s2 = self.d2(s1)
            #print(s2.size())
            s3 = self.d3(s2)
            #print(s3.size())
            s4 = self.d4(s3)
            #print(s4.size())
            scale_1.append(s1)
            scale_2.append(s2)
            scale_3.append(s3)
            scale_4.append(s4)
        scale_1 = torch.cat(scale_1,dim=1)
        #print(scale_1.size())
        scale_2 = torch.cat(scale_2,dim=1)
        scale_3 = torch.cat(scale_3,dim=1)
        scale_4 = torch.cat(scale_4,dim=1)
        scale_1 = self.relu(self.conv1(scale_1))
        scale_2 = self.relu(self.conv2(scale_2))
        scale_3 = self.relu(self.conv3(scale_3))
        scale_4 = self.relu(self.conv4(scale_4))

        return scale_1,scale_2,scale_3,scale_4
"""调节堆栈---end"""


